package top.lyc829.leetcode.T974;

import java.util.HashMap;
import java.util.Map;

public class Solution {

    /**
     * 暴力法
     *
     * @param A
     * @param K
     * @return
     */
    public int subarraysDivByK(int[] A, int K) {

        int res = 0;
        int sum;

        for (int i = 0; i < A.length; i++) {
            sum = 0;
            for (int j = i; j < A.length; j++) {
                sum += A[j];
                if (sum % K == 0) {
                    res++;
                }
            }
        }
        return res;
    }

    /**
     * 前缀和解法
     *
     * @param A
     * @param K
     * @return
     */
    public int subarraysDivByK2(int[] A, int K) {

        // 1, 2, 3
        // 0, 1, 3, 6
        int res = 0;
        int len = A.length;
        int[] tmp = new int[len + 1];
        tmp[0] = 0;
        for (int i = 0; i < len; i++) {
            tmp[i + 1] = tmp[i] + A[i];
        }

        for (int left = 0; left < len; left++) {
            for (int right = left; right < len; right++) {
                if ((tmp[right + 1] - tmp[left]) % K == 0) {
                    res += 1;
                }
            }
        }

        return res;

    }

    /**
     * 前缀和+哈希表
     * 由于只关心次数，不关心具体的解，可以用哈希表记录前缀和被整除的次数
     * 对于前缀和，(sum[j] - sum[i]) % k == 0 -> sum[j] % k = sum[i] % k
     * 统计键对应的值
     *
     * @param A
     * @param K
     * @return
     */
    public int subarraysDivByK3(int[] A, int K) {
        int res = 0;

        Map<Integer, Integer> map = new HashMap<>();
        map.put(0, 1);

        int sum = 0;
        for (int item : A) {
            sum = sum + item;
            int tmp = (sum % K + K) % K;
            int x = map.getOrDefault(tmp, 0);
            res += x;
            map.put(tmp, x + 1);
        }

        return res;
    }


    public static void main(String[] args) {
        Solution s = new Solution();
        int res = s.subarraysDivByK3(new int[]{4, 5, 0, -2, -3, 1}, 5);
        System.out.println(res);
    }
}
